package com.zentrale1.noodlemail.classifiers;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Set;

import com.zentrale1.noodlemail.Mail;

/*
 * The MIT License
 * 
 * Copyright (c) 2010 Andreas Fleig
 * 
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 * 
 * The above copyright notice and this permission notice shall be included in
 * all copies or substantial portions of the Software.
 * 
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 * THE SOFTWARE.
 *
 */


/**
 * Naive Bayes Classifier
 * @author afleig
 *
 */
public class NoodleBayes extends NoodleClassifier 
{	
	// maps category -> word -> occurrences
	private HashMap<String, HashMap<String, Integer>> wordMap;
	
	// maps category -> number of documents
	private HashMap<String, Integer> docMap;
	
	static final double MIN_PROB = 0.01d;
	
	
	public NoodleBayes(String[] folderNames)
	{
		super(folderNames);
		
		wordMap = new HashMap<String, HashMap<String, Integer>>();
		docMap = new HashMap<String, Integer>();
	}
	
	
	public String getName()
	{
		return "Internal Naive Bayes Classifier";
	}

	
	@Override
	protected Result[] _classify(Mail mail) 
	{
		String[] bag = getWordBag(mail);		
		String[] categories = getCategories();
		Result[] results = new Result[categories.length];
		
		// number of all documents
		int allDocs = 0;
		for (String cat : categories) {
			allDocs = allDocs + docMap.get(cat);
		}
		
		// number of all words
		int allWords = 0;
		HashMap<String, Integer> catWords = new HashMap<String, Integer>();
		for (String cat : categories) {
			HashMap<String, Integer> words = wordMap.get(cat);
			Set<String> keys = words.keySet();
			allWords = allWords + keys.size();
			catWords.put(cat, keys.size());
		}
		
		// for each category, get 
		// * pCat: probability that a mail is in category
		// * pCat_Doc: probability of category given a mail (this is the
		//   end result)
		for (int ci=0; ci<categories.length; ci++) {
			String cat = categories[ci];
			HashMap<String, Integer> words = wordMap.get(cat);
			double pCat = new Double(docMap.get(cat)) / new Double(allDocs);
			if (pCat < MIN_PROB)
				pCat = MIN_PROB;
			int catSize = catWords.get(cat);
			double pCat_Word[] = new double[bag.length];
			
			// for each word in mail, get
			// * catOccurrences: number of occurrences of word in category
			// * allOccurrences: number of occurrences in all categories
			// * pWord_Cat: probability of word, given the category
			// * pWord: probability of word
			// * pCat_Word[i]: probability of category given a word
			
			for (int i=0; i<bag.length; i++) {
				int catOccurrences = 0;
				if (words.containsKey(bag[i])) {
					catOccurrences = words.get(bag[i]);
				}
				
				double pWord_Cat = new Double(catOccurrences) / 
					new Double(catSize);
				
				if (pWord_Cat < MIN_PROB)
					pWord_Cat = MIN_PROB;
				
				int allOccurrences = 0;
				for (String c : categories) {
					HashMap<String, Integer> cWords = wordMap.get(c);
					int cO = 0;
					if (cWords.containsKey(bag[i]))
						cO = cWords.get(bag[i]);
					allOccurrences = allOccurrences + cO;
				}
				double pWord = new Double(allOccurrences) / new Double(allWords);
				if (pWord < MIN_PROB)
					pWord = MIN_PROB;
				// bayes!
				pCat_Word[i] = (pWord_Cat * pCat) / pWord;
			}
			
			// pCat_Doc = mean(pCat_Word[i] for all i)
			double pCat_Doc = 1;
			for (double p : pCat_Word) {
				pCat_Doc = pCat_Doc * p;
			}
			pCat_Doc = pCat_Doc / pCat_Word.length;
			
			results[ci] = new Result(cat, pCat_Doc);
		}
		
		return normalize(results);
	}
	
	
	/**
	 * Normalizes the results, so that they add to 1
	 * @param results
	 * @return
	 */
	private Result[] normalize(Result[] results)
	{
		double a = 0;
		for (Result item : results) {
			a += item.getConfidence();
		}
		a = 1.0f / a;
		for (Result item : results) {
			double c = item.getConfidence() * a;
			item.setConfidence(c);
		}

		return results;
	}

	
	@Override
	protected void _train(Mail mail, String category) 
	{
		addCategory(category);
		incrementDocCount(category);
		
		String[] subjectBag = splitString(mail.getSubject());
		addBag(category, subjectBag);
		
		String[] bodyBag = splitString(mail.getTextBody());
		addBag(category, bodyBag);
	}
	
	
	/**
	 * Adds all words from an array to the word list
	 * @param category category to add the words to
	 * @param bag bag of words
	 */
	private void addBag(String category, String[] bag)
	{
		HashMap<String, Integer> words = wordMap.get(category);
		for (String wrd : bag) {
			String word = wrd.toLowerCase();
			if (!word.equals("")) {
				Integer wordCount = 0;
				if (words.containsKey(word))
					wordCount = words.get(word);
				wordCount++;
				words.put(word, wordCount);
			}
		}
	}
	
	
	private String[] splitString(String input)
	{
		return input.split("[\t\r\n.,;\"\' ]");
	}
	
	
	private String[] getWordBag(Mail mail)
	{
		String[] subjectBag = splitString(mail.getSubject());
		String[] bodyBag = splitString(mail.getTextBody());
		ArrayList<String> bag = new ArrayList<String>();
		
		for (String s : subjectBag) {
			if (!s.equals(""))
				bag.add(s.toLowerCase());
		}
		for (String s : bodyBag) {
			if (!s.equals(""))
				bag.add(s.toLowerCase());
		}
		return bag.toArray(new String[bag.size()]);
	}
	
	
	private void addCategory(String cat)
	{
		if (!wordMap.containsKey(cat)) {
			HashMap<String, Integer> words = new HashMap<String, Integer>();
			wordMap.put(cat, words);
		}
	}
	
	
	private void incrementDocCount(String category)
	{
		Integer docCount = 0;
		if (docMap.containsKey(category))
			docCount = docMap.get(category);
		docCount++;
		docMap.put(category, docCount);
	}
	
	
	private String[] getCategories()
	{
		Set<String> keys = docMap.keySet();
		String[] categories = new String[keys.size()];
		
		int i=0;
		for (String key : keys) {
			categories[i] = key;
			i++;
		}
		
		return categories;
	}

}
